#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 7.PyTorch池化层和归一化层.py
# @Author: Richard Chiming Xu
# @Date  : 2021/11/9
# @Desc  :

import torch
from torch import nn
'''
    N: batch
    C: channel
    H: high
    W: weight
    输入维度 (N, C, H_in, W_in)
    输出维度 (N, C, H_out, W_out)
'''
max_pool = nn.MaxPool2d(2)
mean_pool = nn.AvgPool2d(3)

input = torch.randn(20, 16, 30, 24)
max_output = max_pool(input)
mean_output = mean_pool(input)
print('输入数据的维度: {}'.format(input.shape))
print('最大池化的维度: {}'.format(max_output.shape))
print('平均池化的维度: {}'.format(mean_output.shape))


'''
    input = N * H * W
'''
batch_norm = nn.BatchNorm2d(16)
print('batch方向做归一化: {}'.format(batch_norm(input).shape))

'''
    input = C * H * W
'''
layer_norm = nn.LayerNorm([16,30,24])
print('channel方向做归一化: {}'.format(layer_norm(input).shape))
